#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Quick Hook and Test
"""

import frida
import sys
import time
import requests
import io

# Fix encoding
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')

print("\n" + "="*80)
print("Quick Hook Test - Shumei Encryption Capture")
print("="*80 + "\n")

# Connect and attach
device = frida.get_usb_device()
print(f"[1] Connected to: {device.name}")

session = device.attach(2586)
print(f"[2] Attached to PID 2586")

# Load simplified hook
hook_code = """
console.log("\\n[*] Hook loaded successfully!\\n");

// Hook all encrypt methods in libsmsdk.so
var smsdk = Module.findModuleByName("libsmsdk.so");
if (smsdk) {
    console.log("[+] Found libsmsdk.so at:", smsdk.base);
    
    // Try to hook common encryption functions
    var exports = smsdk.enumerateExports();
    console.log("[+] Total exports:", exports.length);
    
    exports.forEach(function(exp) {
        if (exp.name.toLowerCase().indexOf('encrypt') !== -1 || 
            exp.name.toLowerCase().indexOf('aes') !== -1) {
            console.log("[!] Found crypto function:", exp.name);
        }
    });
}

// Hook Java crypto
Java.perform(function() {
    console.log("\\n[+] Hooking Java crypto...");
    
    // Hook AES encryption
    var Cipher = Java.use("javax.crypto.Cipher");
    Cipher.doFinal.overload('[B').implementation = function(data) {
        console.log("\\n========== AES Encryption Detected ==========");
        console.log("[*] Algorithm:", this.getAlgorithm());
        console.log("[*] Input length:", data.length);
        console.log("[*] Input (first 32 bytes):", hexdump(data.slice(0, 32)));
        
        var result = this.doFinal(data);
        console.log("[*] Output length:", result.length);
        console.log("[*] Output (first 32 bytes):", hexdump(result.slice(0, 32)));
        console.log("===========================================\\n");
        
        return result;
    };
    
    console.log("[OK] Hooks ready!\\n");
});

function hexdump(arr) {
    if (!arr || arr.length === 0) return "";
    var result = "";
    for (var i = 0; i < arr.length && i < 32; i++) {
        result += ("0" + (arr[i] & 0xFF).toString(16)).slice(-2) + " ";
    }
    return result;
}
"""

script = session.create_script(hook_code)

captured_data = []

def on_message(message, data):
    try:
        if message['type'] == 'send':
            print(f"[Frida] {message['payload']}")
        elif message['type'] == 'error':
            print(f"[ERROR] {message['stack']}")
    except Exception as e:
        print(f"[Message Error] {e}")

script.on('message', on_message)
script.load()
print("[3] Hook script loaded\n")

time.sleep(2)

# Call API
print("="*80)
print("[4] Calling API to trigger encryption...")
print("="*80 + "\n")

try:
    response = requests.get("http://120.27.155.222:9999/api/message/send-test", timeout=10)
    print(f"[OK] API Response: {response.status_code}\n")
    data = response.json()
    print(f"dev (base64): {data['dev'][:80]}...")
except Exception as e:
    print(f"[ERROR] {e}")

print("\n[5] Waiting for hook output (10 seconds)...")
time.sleep(10)

print("\n" + "="*80)
print("[6] Test complete! Press Ctrl+C to exit...")
print("="*80)

try:
    while True:
        time.sleep(1)
except KeyboardInterrupt:
    print("\n[+] Exiting...\n")
    session.detach()


